package com.aiwiown.snackmq.cluster.metadata;

import com.aiwiown.snackmq.cluster.config.ClusterConfig;
import com.aiwiown.snackmq.cluster.raft.SnackMQRaftNodeConfig;
import com.aiwiown.snackmq.cluster.raft.entity.SnackMQRaftLogEntry;
import com.aiwiown.snackmq.cluster.raft.manager.SnackMQRaftNodeManager;
import com.aiwiown.snackmq.cluster.routing.Binding;
import com.aiwiown.snackmq.cluster.routing.Exchange;
import com.aiwiown.snackmq.cluster.routing.exchange.DirectExchange;
import com.aiwiown.snackmq.cluster.routing.exchange.FanoutExchange;
import com.aiwiown.snackmq.cluster.routing.exchange.TopicExchange;
import com.aiwiown.snackmq.cluster.serialization.ExchangeTypeAdapter;
import com.aiwiown.snackmq.cluster.util.BrokerIdUtils;
import com.aiwiown.snackmq.common.auth.Permission;
import com.aiwiown.snackmq.common.auth.User;
import com.aiwiown.snackmq.common.cluster.BrokerId;
import com.aiwiown.snackmq.common.cluster.MetadataChangeEvent;
import com.aiwiown.snackmq.common.cluster.MetadataListener;
import com.aiwiown.snackmq.common.protocol.MetadataOperation;
import com.aiwiown.snackmq.common.protocol.OperationType;
import com.aiwiown.snackmq.common.protocol.TopicPartition;
import com.aiwiown.snackmq.common.routing.ExchangeType;
import com.alipay.sofa.jraft.Closure;
import com.alipay.sofa.jraft.Iterator;
import com.alipay.sofa.jraft.Status;
import com.alipay.sofa.jraft.conf.Configuration;
import com.alipay.sofa.jraft.core.StateMachineAdapter;
import com.alipay.sofa.jraft.entity.PeerId;
import com.alipay.sofa.jraft.error.RaftError;
import com.alipay.sofa.jraft.rpc.RpcServer;
import com.alipay.sofa.jraft.storage.snapshot.SnapshotReader;
import com.alipay.sofa.jraft.storage.snapshot.SnapshotWriter;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

@Slf4j
public class MetadataManager extends StateMachineAdapter {

    public static final String METADATA_GROUP_ID = "snackmq-meta-group";
    private final SnackMQRaftNodeManager raftNodeManager;
    // --- 【新增】元数据快照支持 ---
    private static final String SNAPSHOT_FILE_NAME = "metadata.snapshot";
    private final Gson gson = new GsonBuilder()
            .registerTypeAdapter(Exchange.class, new ExchangeTypeAdapter())
            .create();
    private final AtomicLong leaderTerm = new AtomicLong(-1);
    private final Rebalancer rebalancer = new Rebalancer();
    private ScheduledExecutorService partitionLeaderCoordinator;

    // --- 元数据核心存储 ---
    private final Map<String, TopicMetadata> topics = new ConcurrentHashMap<>();
    private final Map<String, Exchange> exchanges = new ConcurrentHashMap<>();
    // 注意：此列表仅用于初始化，不应作为集群实时状态的依据。
    private final List<BrokerId> initialBrokers = new CopyOnWriteArrayList<>();
    // --- 【新增】RBAC 元数据核心存储 ---
    private final ConcurrentMap<String, User> users = new ConcurrentHashMap<>();
    private final ConcurrentMap<String, List<Permission>> rolePermissions = new ConcurrentHashMap<>();

    // --- 事件监听器 ---
    private final List<MetadataListener> listeners = new CopyOnWriteArrayList<>();

    // --- 新增：集群配置 ---
    private final ClusterConfig clusterConfig;

    /**
     * ======================= 【核心调整：简化构造函数】 =======================
     * 这是现在唯一的、推荐的构造函数。
     * 它清晰地表明 MetadataManager 依赖一个类型安全的 ClusterConfig 对象。
     *
     * @param config        Raft 节点配置
     * @param rpcServer     RPC 服务器
     * @param clusterConfig 集群配置
     */
    public MetadataManager(SnackMQRaftNodeConfig config, RpcServer rpcServer, ClusterConfig clusterConfig) {
        this.clusterConfig = Objects.requireNonNull(clusterConfig, "clusterConfig cannot be null");
        config.setGroupId(METADATA_GROUP_ID);
        this.raftNodeManager = new SnackMQRaftNodeManager(config, this, rpcServer, this.clusterConfig);
        // 【修复】调用 initializeInitialBrokers 时，它会使用新的 getBusinessPortForHost 逻辑
        initializeInitialBrokers(config.getInitialConf());
        exchanges.put(Exchange.DEFAULT_DIRECT_EXCHANGE, new DirectExchange(Exchange.DEFAULT_DIRECT_EXCHANGE));
    }


    /**
     * 向后兼容的构造函数
     *
     * @param config    Raft 节点配置
     * @param rpcServer RPC 服务器
     */
    @Deprecated
    public MetadataManager(SnackMQRaftNodeConfig config, RpcServer rpcServer) {
        this(config, rpcServer, new ClusterConfig());
    }

    /**
     * 向后兼容的构造函数，接受 BrokerConfig 对象
     *
     * @param config    Raft 节点配置
     * @param rpcServer RPC 服务器
     * @param brokerConfig Broker 配置对象
     */
    @Deprecated
    public MetadataManager(SnackMQRaftNodeConfig config, RpcServer rpcServer, Object brokerConfig) {
        this(config, rpcServer, ClusterConfig.fromBrokerConfig(brokerConfig));
    }

    private void initializeInitialBrokers(String initialConf) {
        if (initialConf == null || initialConf.isEmpty()) {
            log.warn("Raft initial configuration is empty, cannot initialize brokers.");
            return;
        }
        for (String peerStr : initialConf.split(",")) {
            PeerId peer = new PeerId();
            if (peer.parse(peerStr)) {
                // 【核心修复】使用新的、可靠的方式获取业务端口
                int businessPort = getBusinessPortForHost(peer.getIp(), peer.getPort());
                initialBrokers.add(BrokerIdUtils.fromRaftPeerWithBusinessPort(peer, businessPort));
            } else {
                log.error("Failed to parse peer from: {}", peerStr);
            }
        }
        log.info("初始化 Broker 配置列表: {}", initialBrokers);
    }

    /**
     * 【核心修复】根据主机地址和 Raft 端口获取业务端口。
     * 这个方法现在变得更加健壮和明确。
     *
     * @param host 主机地址
     * @param raftPort Raft 端口
     * @return 业务端口
     */
    private int getBusinessPortForHost(String host, int raftPort) {
        // 1. 优先从显式映射中查找
        if (clusterConfig != null && clusterConfig.getRaftAddressToBusinessPortMap() != null) {
            String raftAddress = host + ":" + raftPort;
            Integer businessPort = clusterConfig.getRaftAddressToBusinessPortMap().get(raftAddress);
            if (businessPort != null) {
                return businessPort;
            }
        }
        // 2. 如果映射不存在或未找到，检查是否是当前节点
        if (clusterConfig != null && host.equals(clusterConfig.getHost()) && raftPort == clusterConfig.getRaftPort()) {
            return clusterConfig.getBusinessPort();
        }
        // 3. 如果所有方法都失败，则回退到旧的约定，并记录一个严重警告。
        log.warn("无法在地址映射中找到主机 {}:{} 的业务端口，将回退到 'raftPort - 1' 的约定。这在集群环境中可能导致错误！", host, raftPort);
        return raftPort - 1;
    }

    /**
     * 添加元数据监听器。
     *
     * @param listener 元数据监听器
     */
    public void addMetadataListener(MetadataListener listener) {
        this.listeners.add(listener);
    }

    /**
     * 通知所有监听器。
     *
     * @param event 元数据变更事件
     */
    private void notifyListeners(MetadataChangeEvent event) {
        for (MetadataListener listener : listeners) {
            try {
                listener.onMetadataChanged(event);
            } catch (Exception e) {
                log.error("执行元数据监听器时出错", e);
            }
        }
    }

    /**
     * 当 Raft 配置提交时调用，通常用于处理集群成员变更。
     *
     * @param conf Raft 配置
     */
    @Override
    public void onConfigurationCommitted(Configuration conf) {
        log.info("Raft configuration changed: {}", conf.getPeers());
        // 当集群成员发生变化（有节点加入或离开）时，也触发一次再均衡。
        triggerRebalance();
    }

    /**
     * 触发再均衡操作，这是修复故障转移问题的关键方法。
     * 它现在基于 Raft 集群的实时存活节点列表来计算分区分配，而不是一个静态的列表。
     */
    private void triggerRebalance() {
        // 确保只有 Leader 才能执行此操作
        if (!isLeader()) {
            return;
        }

        // 1. 从 Raft Node 获取当前所有 *存活* 的对等节点 (Live Peers)。这才是决定分区分配的权威数据源。
        //    这个改动是解决问题的关键。
        List<PeerId> alivePeers = this.raftNodeManager.getAlivePeers();
        if (alivePeers.isEmpty()) {
            log.warn("无法触发再均衡，在 Raft 集群中没有发现存活的 Broker。");
            return;
        }

        // 2. 将存活的 Raft 节点 (host:raft_port) 转换为 BrokerId (host:business_port)。
        List<BrokerId> currentLiveBrokers = alivePeers.stream()
                .map(peer -> BrokerIdUtils.fromRaftPeerWithBusinessPort(peer, getBusinessPortForHost(peer.getIp(), peer.getPort())))
                .collect(Collectors.toList());

        log.info("触发再均衡，当前存活的 Broker 列表: {}", currentLiveBrokers);

        // 3. 使用 *仅仅* 存活的 Broker 列表来计算新的分区分配方案。
        Map<TopicPartition, List<BrokerId>> newAssignment = rebalancer.calculateNewAssignments(this.topics, currentLiveBrokers);

        if (newAssignment.isEmpty() && !this.topics.isEmpty()) {
            log.warn("再均衡计算出了一个空的分配方案，可能存在问题，跳过更新。当前存活节点: {}", currentLiveBrokers);
            return;
        }

        // 4. 提议元数据变更
        MetadataOperation updateOp = MetadataOperation.newUpdatePartitionAssignmentOperation(newAssignment);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, updateOp);
        raftNodeManager.apply(logEntry, status -> {
            if (status.isOk()) {
                log.info("成功提议了再均衡操作。");
            } else {
                log.error("提议再均衡操作失败: {}", status);
            }
        });
    }

    /**
     * 当有新的日志条目应用到状态机时调用。
     *
     * @param iter 日志条目迭代器
     */
    @Override
    public void onApply(Iterator iter) {
        while (iter.hasNext()) {
            SnackMQRaftLogEntry logEntry = null;
            MetadataOperation operation = null; // 【诊断增强】将 operation 声明移到外部，以便在 catch 块中访问
            Closure closure = iter.done();
            try {
                log.debug("Applying Raft log at index {}", iter.getIndex());
                logEntry = SnackMQRaftLogEntry.fromBytes(iter.getData().array());
                if (logEntry != null && logEntry.getOperationType() == OperationType.METADATA) {
                    log.debug("Deserializing MetadataOperation from log entry...");
                    operation = logEntry.getData(MetadataOperation.class);
                    log.debug("Successfully deserialized operation: {}", operation != null ? operation.getOperationType() : "null");

                    applyOperation(operation);
                } else {
                    log.warn("在元数据状态机中收到了非元数据类型的日志: {}", logEntry != null ? logEntry.getOperationType() : "null");
                }
                if (closure != null) {
                    closure.run(Status.OK());
                }
            } catch (Exception e) {
                // 【诊断增强】提供更详细的错误日志，明确指出失败的操作类型和 Raft 索引
                String operationTypeInfo = (operation != null) ? String.valueOf(operation.getOperationType()) :
                        (logEntry != null) ? "LogEntry(type=" + logEntry.getOperationType() + ")" : "Unknown";

                log.error("应用元数据操作失败。操作类型 [{}], Raft索引 [{}].",
                        operationTypeInfo, iter.getIndex(), e);

                if (closure != null) {
                    closure.run(new Status(RaftError.EINTERNAL, "应用元数据操作失败。"));
                }
            }
            iter.next();
        }
    }

    /**
     * 应用元数据操作。
     *
     * @param operation 元数据操作
     */
    private void applyOperation(MetadataOperation operation) {
        if (operation == null) return;
        if (operation.getType() != OperationType.METADATA) {
            log.warn("只支持 METADATA 类型的操作，实际收到: {}", operation.getType());
            return;
        }
        switch (operation.getOperationType()) {
            case CREATE_TOPIC:
                applyCreateTopic(operation);
                break;
            case DELETE_TOPIC:
                applyDeleteTopic(operation);
                break;
            case UPDATE_PARTITION_ASSIGNMENT:
                applyUpdatePartitionAssignment(operation);
                break;
            case UPDATE_LEADER:
                applyUpdateLeader(operation);
                break;
            case DECLARE_EXCHANGE:
                applyDeclareExchange(operation);
                break;
            case BIND_QUEUE:
                applyBindQueue(operation);
                break;
            case UNBIND_QUEUE:
                applyUnbindQueue(operation);
                break;
            case DELETE_EXCHANGE:
                applyDeleteExchange(operation);
                break;
            case CREATE_USER:
                applyCreateUser(operation);
                break;
            case SET_ROLE_PERMISSIONS:
                applySetRolePermissions(operation);
                break;
            case DELETE_USER:
                applyDeleteUser(operation);
                break;
            case DELETE_ROLE:
                applyDeleteRole(operation);
                break;
            default:
                log.warn("不支持的元数据操作类型: {}", operation.getOperationType());
        }
    }

    /**
     * 【新增】一个私有内部类，用于封装所有需要持久化到快照的元数据。
     * 这样做可以简化序列化和反序列化的过程。
     */
    private static class MetadataSnapshot {
        private Map<String, TopicMetadata> topics;
        private Map<String, Exchange> exchanges;
        private Map<String, User> users;
        private Map<String, List<Permission>> rolePermissions;
    }

    @Override
    public void onSnapshotSave(SnapshotWriter writer, Closure done) {
        // 【新增】实现元数据快照的保存逻辑。这是解决重启后状态丢失问题的关键。
        log.info("开始保存元数据快照...");
        MetadataSnapshot snapshot = new MetadataSnapshot();
        snapshot.topics = new HashMap<>(this.topics);
        snapshot.exchanges = new HashMap<>(this.exchanges);
        snapshot.users = new HashMap<>(this.users);
        snapshot.rolePermissions = new HashMap<>(this.rolePermissions);

        String snapshotPath = writer.getPath();
        File snapshotFile = new File(snapshotPath, SNAPSHOT_FILE_NAME);

        try {
            String jsonState = gson.toJson(snapshot);
            Files.write(snapshotFile.toPath(), jsonState.getBytes());
            if (writer.addFile(SNAPSHOT_FILE_NAME)) {
                log.info("成功保存元数据快照，包含 {} 个主题, {} 个交换机, {} 个用户, {} 个角色。",
                        snapshot.topics.size(), snapshot.exchanges.size(), snapshot.users.size(), snapshot.rolePermissions.size());
                done.run(Status.OK());
            } else {
                log.error("无法将快照文件 {} 添加到 writer。", snapshotFile.getAbsolutePath());
                done.run(new Status(RaftError.EIO, "无法添加快照文件。"));
            }
        } catch (IOException e) {
            log.error("保存元数据快照失败。", e);
            done.run(new Status(RaftError.EIO, "保存元数据快照时发生 I/O 错误。"));
        }
    }

    @Override
    public boolean onSnapshotLoad(SnapshotReader reader) {
        // 【新增】实现从快照加载元数据的逻辑。
        if (reader.listFiles() == null || !reader.listFiles().contains(SNAPSHOT_FILE_NAME)) {
            log.warn("在快照中找不到元数据文件 {}，将从空状态开始。", SNAPSHOT_FILE_NAME);
            return true;
        }

        File snapshotFile = new File(reader.getPath(), SNAPSHOT_FILE_NAME);
        try {
            String jsonState = new String(Files.readAllBytes(snapshotFile.toPath()));
            MetadataSnapshot snapshot = gson.fromJson(jsonState, MetadataSnapshot.class);

            // 原子地替换当前状态
            this.topics.clear();
            this.topics.putAll(snapshot.topics);
            this.exchanges.clear();
            this.exchanges.putAll(snapshot.exchanges);
            this.users.clear();
            this.users.putAll(snapshot.users);
            this.rolePermissions.clear();
            this.rolePermissions.putAll(snapshot.rolePermissions);

            log.info("成功从快照加载元数据，包含 {} 个主题, {} 个交换机, {} 个用户, {} 个角色。",
                    this.topics.size(), this.exchanges.size(), this.users.size(), this.rolePermissions.size());
            return true;
        } catch (Exception e) {
            log.error("从快照加载元数据失败。", e);
            return false;
        }
    }

    /**
     * 【新增】通过 Raft 提议删除一个用户。
     */
    public void deleteUser(String username, Closure closure) {
        MetadataOperation op = MetadataOperation.newDeleteUserOperation(username);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 【新增】通过 Raft 提议删除一个角色。
     */
    public void deleteRole(String roleName, Closure closure) {
        MetadataOperation op = MetadataOperation.newDeleteRoleOperation(roleName);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 【新增】应用删除用户操作。
     */
    private void applyDeleteUser(MetadataOperation op) {
        String username = op.getUsernameToDelete();
        if (username != null && users.remove(username) != null) {
            log.info("Applied: Delete User '{}'", username);
        }
    }

    /**
     * 【新增】应用删除角色操作。
     * 这也会从所有拥有该角色的用户中移除此角色。
     */
    private void applyDeleteRole(MetadataOperation op) {
        String roleName = op.getRoleNameToDelete();
        if (roleName != null && rolePermissions.remove(roleName) != null) {
            log.info("Applied: Delete Role '{}'. Now removing it from all users.", roleName);
            // 从所有用户中移除此角色
            users.values().forEach(user -> {
                if (user.getRoles() != null && user.getRoles().contains(roleName)) {
                    // 创建一个新的、可修改的列表
                    List<String> updatedRoles = new ArrayList<>(user.getRoles());
                    updatedRoles.remove(roleName);
                    user.setRoles(updatedRoles);
                    log.info("Removed deleted role '{}' from user '{}'", roleName, user.getUsername());
                }
            });
        }
    }

    /**
     * 【新增】应用创建用户操作。
     */
    private void applyCreateUser(MetadataOperation op) {
        User user = op.getUser();
        if (user != null && !users.containsKey(user.getUsername())) {
            users.put(user.getUsername(), user);
            log.info("Applied: Create User '{}' with roles {}", user.getUsername(), user.getRoles());
        }
    }

    /**
     * 【新增】应用设置角色权限操作。
     */
    private void applySetRolePermissions(MetadataOperation op) {
        String roleName = op.getRoleName();
        List<Permission> permissions = op.getPermissions();
        if (roleName != null && permissions != null) {
            rolePermissions.put(roleName, new CopyOnWriteArrayList<>(permissions));
            log.info("Applied: Set {} permissions for role '{}'", permissions.size(), roleName);
        }
    }

    /**
     * 根据用户名获取用户信息。
     * @param username 用户名
     * @return User 对象，如果不存在则返回 null
     */
    public User getUser(String username) {
        User original = users.get(username);
        if (original == null) {
            return null;
        }
        // --- 【最终修复】: 返回一个 User 对象的安全副本，而不是直接引用。---
        // 这是为了防止其他代码（例如，为了安全而清除密码的API层）意外地修改了内存中缓存的原始 User 对象。
        // 直接返回引用是导致密码在运行时变为 null 的根本原因。
        // 我们假设 User 类有一个接受 username, password, roles 的构造函数。
        return new User(original.getUsername(), original.getPassword(), new ArrayList<>(original.getRoles()));
    }

    /**
     * 获取指定角色的所有权限。
     * @param roleName 角色名
     * @return 权限列表，如果角色不存在则返回 null
     */
    public List<Permission> getPermissionsForRole(String roleName) {
        return rolePermissions.get(roleName);
    }

    /**
     * 通过 Raft 提议创建一个新用户。
     * @param user 要创建的用户对象
     * @param closure 操作完成后的回调
     */
    public void createUser(User user, Closure closure) {
        MetadataOperation op = MetadataOperation.newCreateUserOperation(user);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 通过 Raft 提议设置或更新一个角色的权限。
     * @param roleName 角色名
     * @param permissions 该角色的权限列表
     * @param closure 操作完成后的回调
     */
    public void setRolePermissions(String roleName, List<Permission> permissions, Closure closure) {
        MetadataOperation op = MetadataOperation.newSetRolePermissionsOperation(roleName, permissions);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 【新增】获取所有主题的元数据列表。
     * @return 一个包含所有 TopicMetadata 的列表。
     */
    public List<TopicMetadata> getAllTopics() {
        return new ArrayList<>(this.topics.values());
    }

    /**
     * 【新增】获取所有交换机的元数据列表。
     * @return 一个包含所有 Exchange 对象的列表。
     */
    public List<Exchange> getAllExchanges() {
        return new ArrayList<>(this.exchanges.values());
    }

    /**
     * 【新增】获取当前所有存活的 Broker 节点列表。
     * 这是获取集群成员的权威方法，用于控制台展示。
     * @return 一个包含所有存活 BrokerId 的列表。
     */
    public List<BrokerId> getAllBrokers() {
        List<PeerId> alivePeers = this.raftNodeManager.getAlivePeers();
        if (alivePeers == null) {
            return Collections.emptyList();
        }
        return alivePeers.stream()
                .map(peer -> BrokerIdUtils.fromRaftPeerWithBusinessPort(peer, getBusinessPortForHost(peer.getIp(), peer.getPort())))
                .collect(Collectors.toList());
    }

    /**
     * 【新增】获取所有已配置的用户列表。
     * @return 用户列表
     */
    public List<User> getAllUsers() {
        // --- 【最终修复】: 返回一个包含 User 对象安全副本的新列表。---
        // 这确保了调用者无法修改缓存中的原始 User 对象。
        return this.users.values().stream()
                .map(original -> new User(original.getUsername(), original.getPassword(), new ArrayList<>(original.getRoles())))
                .collect(Collectors.toList());
    }

    /**
     * 【新增】获取所有角色及其权限的映射。
     * @return 角色名到权限列表的 Map
     */
    public Map<String, List<Permission>> getAllRolePermissions() {
        // 返回一个副本以确保线程安全和封装性
        return new ConcurrentHashMap<>(this.rolePermissions);
    }
    /**
     * 应用创建主题操作，并触发再均衡。
     *
     * @param op 创建主题操作
     */
    private void applyCreateTopic(MetadataOperation op) {
        String topicName = op.getTopicName();
        if (topics.containsKey(topicName)) {
            log.warn("主题 {} 已存在。", topicName);
            return;
        }
        log.info("收到创建主题 {} (分区: {}, 副本: {}) 的请求，将触发一次再均衡来分配其分区。",
                topicName, op.getPartitions(), op.getReplicationFactor());
        TopicMetadata newTopic = new TopicMetadata(topicName);
        // 【核心修改】将副本因子存入 Topic 元数据
        newTopic.setReplicationFactor(op.getReplicationFactor());
        for (int i = 0; i < op.getPartitions(); i++) {
            newTopic.setPartitionAssignment(i, new TopicMetadata.PartitionAssignment(new ArrayList<>(), null));
        }
        this.topics.put(topicName, newTopic);
        triggerRebalance();
    }

    /**
     * 应用删除主题操作。
     *
     * @param op 删除主题操作
     */
    private void applyDeleteTopic(MetadataOperation op) {
        String topicName = op.getTopicName();
        if (topics.containsKey(topicName)) {
            topics.remove(topicName);
            log.info("Applied: Delete Topic '{}'", topicName);
            // 触发一次再均衡以清理可能残留的分区分配信息
            triggerRebalance();
        }
    }


    /**
     * 应用更新分区分配的元数据操作。
     * 此方法首先检查新分配是否有效，然后更新内部状态，并最后通知监听器。
     *
     * @param op 包含分区分配更新信息的元数据操作对象。
     */
    private void applyUpdatePartitionAssignment(MetadataOperation op) {
        Map<TopicPartition, List<BrokerId>> newAssignment = op.getPartitionAssignment();
        if (newAssignment == null) return;

        Map<TopicPartition, TopicMetadata.PartitionAssignment> oldAssignments = new HashMap<>();
        topics.values().forEach(tm -> tm.getPartitionAssignments().forEach((pid, pa) ->
                oldAssignments.put(new TopicPartition(tm.getTopicName(), pid), pa)));

        newAssignment.forEach((tp, newReplicas) -> {
            // 【核心修复】: 使用 get() 代替 computeIfAbsent()。
            // 这可以防止在应用一个过时的分区分配日志时，意外地“复活”一个已经被删除的主题。
            TopicMetadata topic = topics.get(tp.getTopic());
            if (topic == null) {
                log.warn("在应用分区分配时，发现主题 '{}' 不存在，将跳过此分区的分配。这可能是由于一个延迟的再均衡操作导致的。", tp.getTopic());
                return; // 在 forEach 中，return 相当于 continue
            }
            TopicMetadata.PartitionAssignment oldAssignment = oldAssignments.get(tp);

            BrokerId leader = null;
            if (oldAssignment != null && newReplicas.contains(oldAssignment.getLeader())) {
                leader = oldAssignment.getLeader();
            } else if (!newReplicas.isEmpty()) {
                leader = newReplicas.get(0);
            }

            topic.setPartitionAssignment(tp.getPartitionId(), new TopicMetadata.PartitionAssignment(newReplicas, leader));
        });
        log.info("已应用新的分区分配方案。");

        // --- 核心修复：在这里发布事件 ---
        Map<String, Object> context = new HashMap<>();
        context.put("newAssignment", newAssignment);
        notifyListeners(new MetadataChangeEvent(MetadataChangeEvent.EventType.PARTITION_ASSIGNMENT_CHANGED, context));
    }

    /**
     * 应用更新Leader的元数据操作。
     * 此方法验证操作的有效性，更新分区的Leader，并通知监听器。
     *
     * @param op 包含Leader更新信息的元数据操作对象。
     */
    private void applyUpdateLeader(MetadataOperation op) {
        TopicPartition partition = op.getTopicPartitionForLeaderUpdate();
        BrokerId newLeader = op.getNewLeader();

        if (partition == null || newLeader == null) {
            log.error("Invalid UPDATE_LEADER operation: {}", op);
            return;
        }

        TopicMetadata topic = topics.get(partition.getTopic());
        if (topic == null) {
            log.warn("Cannot update leader for non-existent topic: {}", partition.getTopic());
            return;
        }

        TopicMetadata.PartitionAssignment assignment = topic.getAssignmentFor(partition.getPartitionId());
        if (assignment == null) {
            log.warn("Cannot update leader for non-existent partition: {}", partition);
            return;
        }

        if (!assignment.getReplicas().contains(newLeader)) {
            log.error("Cannot assign leader {} to partition {} because it is not in the replica set {}.",
                    newLeader, partition, assignment.getReplicas());
            return;
        }

        if (Objects.equals(assignment.getLeader(), newLeader)) {
            log.info("Leader for partition {} is already {}. No change needed.", partition, newLeader);
            return;
        }

        log.info("Applying leader change for partition {}. New leader: {}", partition, newLeader);
        assignment.setLeader(newLeader);

        // --- 核心修复：在这里发布事件 ---
        Map<String, Object> context = new HashMap<>();
        context.put("partition", partition);
        context.put("newLeader", newLeader);
        notifyListeners(new MetadataChangeEvent(MetadataChangeEvent.EventType.LEADER_CHANGED, context));
    }

    /**
     * 应用声明交换机的元数据操作。
     * 此方法检查交换机是否已存在，如果不存在则创建新的交换机。
     *
     * @param op 包含交换机声明信息的元数据操作对象。
     */
    private void applyDeclareExchange(MetadataOperation op) {
        String exchangeName = op.getExchangeName();
        ExchangeType type = op.getExchangeType();

        if (exchanges.containsKey(exchangeName)) {
            log.warn("Exchange {} already exists.", exchangeName);
            return;
        }

        Exchange newExchange;
        switch (type) {
            case FANOUT:
                newExchange = new FanoutExchange(exchangeName);
                break;
            case TOPIC:
                newExchange = new TopicExchange(exchangeName);
                break;
            case DIRECT:
            default:
                newExchange = new DirectExchange(exchangeName);
                break;
        }
        exchanges.put(exchangeName, newExchange);
        log.info("Applied: Declare Exchange '{}' of type {}", exchangeName, type);
    }

    /**
     * 应用绑定队列的元数据操作。
     * 此方法检查交换机是否存在，如果存在则创建新的绑定。
     *
     * @param op 包含绑定队列信息的元数据操作对象。
     */
    private void applyBindQueue(MetadataOperation op) {
        Exchange exchange = exchanges.get(op.getExchangeName());
        if (exchange != null) {
            Binding binding = new Binding(op.getExchangeName(), op.getQueueName(), op.getRoutingKey());
            exchange.addBinding(binding);
            log.info("Applied: Bind Queue '{}' to Exchange '{}' with key '{}'", op.getQueueName(), op.getExchangeName(), op.getRoutingKey());
        } else {
            log.warn("Cannot apply binding, exchange '{}' not found.", op.getExchangeName());
        }
    }

//    /**
//     * 【新增】应用解除队列绑定的元数据操作。
//     * 此方法会从指定的交换机中移除一个绑定。
//     *
//     * @param op 包含解除绑定信息的元数据操作对象。
//     */
//    private void applyUnbindQueue(MetadataOperation op) {
//        Exchange oldExchange = exchanges.get(op.getExchangeName());
//        if (oldExchange != null) {
//            // 【修复】getBindings() 返回的是一个不可修改的列表，直接对其调用 removeIf 或 iterator.remove() 会抛出 UnsupportedOperationException。
//            // 正确的方式是先创建一个可修改的列表副本。
//            final List<Binding> newBindings = new ArrayList<>(oldExchange.getBindings());
//
//            // 使用 removeIf 在可修改的副本上安全地移除目标绑定
//            boolean wasRemoved = newBindings.removeIf(binding ->
//                    Objects.equals(binding.getQueueName(), op.getQueueName()) &&
//                            Objects.equals(binding.getRoutingKey(), op.getRoutingKey())
//            );
//
//            if (wasRemoved) {
//                // 只有在确实删除了一个绑定时，才需要重建和替换 Exchange 对象
//                final Exchange newExchange;
//                switch (oldExchange.getType()) {
//                    case FANOUT:
//                        newExchange = new FanoutExchange(oldExchange.getName());
//                        break;
//                    case TOPIC:
//                        newExchange = new TopicExchange(oldExchange.getName());
//                        break;
//                    case DIRECT:
//                    default:
//                        newExchange = new DirectExchange(oldExchange.getName());
//                        break;
//                }
//                newBindings.forEach(newExchange::addBinding);
//                this.exchanges.put(oldExchange.getName(), newExchange);
//                log.info("Applied: Unbind Queue '{}' from Exchange '{}' with key '{}'", op.getQueueName(), oldExchange.getName(), op.getRoutingKey());
//            } else {
//                // 【诊断增强】新增日志，用于调试未找到匹配绑定的情况
//                log.warn("Attempted to unbind queue, but no matching binding was found. Exchange: '{}', Queue: '{}', Key: '{}'",
//                        op.getExchangeName(), op.getQueueName(), op.getRoutingKey());
//            }
//        } else {
//            log.warn("Attempted to unbind queue from a non-existent exchange: '{}'", op.getExchangeName());
//        }
//    }

    /**
     * 【新增】应用解除队列绑定的元数据操作。
     * 此方法会从指定的交换机中移除一个绑定。
     *
     * @param op 包含解除绑定信息的元数据操作对象。
     */
    private void applyUnbindQueue(MetadataOperation op) {
        Exchange exchange = exchanges.get(op.getExchangeName());
        if (exchange != null) {
            // 【优化】直接在现有的 Exchange 实例上移除绑定。
            // Exchange 的实现类内部使用 CopyOnWriteArrayList，
            // 这使得移除操作本身是线程安全的，并且比重新创建整个 Exchange 对象高效得多。
            Binding bindingToRemove = new Binding(op.getExchangeName(), op.getQueueName(), op.getRoutingKey());
            boolean wasRemoved = exchange.getBindings().contains(bindingToRemove);
            exchange.removeBinding(bindingToRemove);

            if (wasRemoved) {
                log.info("Applied: Unbind Queue '{}' from Exchange '{}' with key '{}'", op.getQueueName(), op.getExchangeName(), op.getRoutingKey());
            } else {
                log.warn("Attempted to unbind queue, but no matching binding was found. Exchange: '{}', Queue: '{}', Key: '{}'",
                        op.getExchangeName(), op.getQueueName(), op.getRoutingKey());
            }
        } else {
            log.warn("Attempted to unbind queue from a non-existent exchange: '{}'", op.getExchangeName());
        }
    }

    /**
     * 【新增】应用删除交换机的元数据操作。
     *
     * @param op 包含要删除的交换机名称的元数据操作。
     */
    private void applyDeleteExchange(MetadataOperation op) {
        String exchangeName = op.getExchangeName();
        if (exchangeName != null && exchanges.containsKey(exchangeName)) {
            // 检查是否是默认交换机，默认交换机不允许删除
            if (Exchange.DEFAULT_DIRECT_EXCHANGE.equals(exchangeName)) {
                log.warn("Attempted to delete the default exchange '{}', which is not allowed.", exchangeName);
                return;
            }
            exchanges.remove(exchangeName);
            log.info("Applied: Delete Exchange '{}'", exchangeName);
        }
    }

    /**
     * 创建一个新主题。
     * 如果主题已存在，则记录警告并返回。
     *
     * @param topicName         主题名称。
     * @param partitions        主题的分区数量。
     * @param replicationFactor 每个分区的副本因子。
     * @param closure           操作完成后的回调。
     */
    public void createTopic(String topicName, int partitions, int replicationFactor, Closure closure) {
        if (this.topics.containsKey(topicName)) {
            log.warn("Topic {} already exists, skipping creation.", topicName);
            closure.run(Status.OK());
            return;
        }
        MetadataOperation op = MetadataOperation.newCreateTopicOperation(topicName, partitions, replicationFactor);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 删除一个主题。
     *
     * @param topicName 主题名称。
     * @param closure   操作完成后的回调。
     */
    public void deleteTopic(String topicName, Closure closure) {
        // 【修复】移除在提议前的存在性检查。所有状态决策都应在状态机内一致地做出，以保证幂等性。
        // 之前的逻辑是，如果 Leader 的本地缓存认为 Topic 不存在，则不会调用 closure，导致客户端请求挂起。
        MetadataOperation op = MetadataOperation.newDeleteTopicOperation(topicName);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 声明一个交换机。
     *
     * @param exchangeName 交换机名称。
     * @param type         交换机类型。
     * @param closure      操作完成后的回调。
     */
    public void declareExchange(String exchangeName, ExchangeType type, Closure closure) {
        MetadataOperation op = MetadataOperation.newDeclareExchangeOperation(exchangeName, type);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 【新增】删除一个交换机。
     *
     * @param exchangeName 要删除的交换机名称。
     * @param closure      操作完成后的回调。
     */
    public void deleteExchange(String exchangeName, Closure closure) {
        // 假设 MetadataOperation.newDeleteExchangeOperation 已被创建
        MetadataOperation op = MetadataOperation.newDeleteExchangeOperation(exchangeName);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 绑定一个队列到交换机。
     *
     * @param queueName    队列名称。
     * @param exchangeName 交换机名称。
     * @param routingKey   绑定的路由键。
     * @param closure      操作完成后的回调。
     */
    public void bindQueue(String queueName, String exchangeName, String routingKey, Closure closure) {
        MetadataOperation op = MetadataOperation.newBindQueueOperation(queueName, exchangeName, routingKey);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 【新增】解除一个队列与交换机的绑定。
     *
     * @param queueName    队列名称。
     * @param exchangeName 交换机名称。
     * @param routingKey   绑定的路由键。
     * @param closure      操作完成后的回调。
     */
    public void unbindQueue(String queueName, String exchangeName, String routingKey, Closure closure) {
        // 假设 MetadataOperation.newUnbindQueueOperation 已被创建
        MetadataOperation op = MetadataOperation.newUnbindQueueOperation(queueName, exchangeName, routingKey);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }
    /**
     * 更新主题分区的Leader。
     * 此方法首先检查节点是否为Leader，然后应用更新操作。
     *
     * @param partition 需要更新Leader的分区。
     * @param newLeader 新的Leader Broker ID。
     * @param closure   操作完成后的回调。
     */
    public void updatePartitionLeader(TopicPartition partition, BrokerId newLeader, Closure closure) {
        if (!isLeader()) {
            closure.run(new Status(RaftError.EPERM, "Not the metadata leader."));
            return;
        }
        MetadataOperation op = MetadataOperation.newUpdateLeaderOperation(partition, newLeader);
        SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.METADATA, op);
        raftNodeManager.apply(logEntry, closure);
    }

    /**
     * 获取指定主题的元数据。
     *
     * @param topicName 主题名称。
     * @return 主题的元数据，如果主题不存在则返回null。
     */
    public TopicMetadata getTopic(String topicName) {
        return topics.get(topicName);
    }

    /**
     * 获取指定交换机的实例。
     *
     * @param exchangeName 交换机名称。
     * @return 交换机实例，如果交换机不存在则返回null。
     */
    public Exchange getExchange(String exchangeName) {
        return exchanges.get(exchangeName);
    }

    /**
     * 根据给定的主题分区，获取其 Leader Broker 的地址。
     * <p>
     * 此方法是线程安全的，因为它从基于 Raft 的、并发安全的数据结构中读取信息。
     *
     * @param partition 需要查询的主题分区。
     * @return 该分区的 Leader Broker 地址 (例如 "localhost:9880")，如果未找到则返回 null。
     */
    public String getPartitionLeaderAddress(TopicPartition partition) {
        if (partition == null) {
            return null;
        }
        TopicMetadata topicMetadata = topics.get(partition.getTopic());
        if (topicMetadata == null) {
            log.warn("Cannot find leader for a non-existent topic: {}", partition.getTopic());
            return null;
        }
        TopicMetadata.PartitionAssignment assignment = topicMetadata.getAssignmentFor(partition.getPartitionId());
        if (assignment == null) {
            log.warn("Cannot find leader for a non-existent partition: {}", partition);
            return null;
        }
        BrokerId leader = assignment.getLeader();
        if (leader == null) {
            log.warn("No leader is currently assigned for partition: {}", partition);
            return null;
        }
        // 【代码修复】: 将 getIp() 修改为 getHost()，与 BrokerId 的定义保持一致。
        return leader.getHost() + ":" + leader.getPort();
    }

    /**
     * 检查指定分区是否分配给指定的 Broker。
     *
     * @param topicPartition 主题分区
     * @param brokerAddress  Broker 地址
     * @return 如果分区分配给该 Broker 则返回 true，否则返回 false
     */
    public boolean isPartitionAssignedToBroker(TopicPartition topicPartition, String brokerAddress) {
        if (topicPartition == null || brokerAddress == null) {
            return false;
        }
        
        TopicMetadata topicMetadata = topics.get(topicPartition.getTopic());
        if (topicMetadata == null) {
            return false;
        }

        TopicMetadata.PartitionAssignment assignment = topicMetadata.getAssignmentFor(topicPartition.getPartitionId());
        if (assignment == null || assignment.getReplicas() == null) {
            return false;
        }

        // 检查该 Broker 是否在副本列表中
        return assignment.getReplicas().stream()
                .anyMatch(brokerId -> (brokerId.getHost() + ":" + brokerId.getPort()).equals(brokerAddress));
    }

    /**
     * 启动 MetadataManager。
     * 如果启动成功，此方法还会调度一个任务来周期性地检查和协调分区的Leader。
     *
     * @param startRpcServer 是否启动RPC服务器。
     * @return 启动是否成功。
     */
    public boolean start(boolean startRpcServer) {
        boolean result = this.raftNodeManager.start(startRpcServer);
        if (result) {
            // 【新增】启动分区 Leader 协调任务
            this.partitionLeaderCoordinator = Executors.newSingleThreadScheduledExecutor(
                    r -> new Thread(r, "partition-leader-coordinator"));
            // 每 5 秒检查一次
            this.partitionLeaderCoordinator.scheduleAtFixedRate(
                    this::coordinatePartitionLeaders, 10, 5, TimeUnit.SECONDS);
        }
        return result;
    }

    /**
     * 关闭 MetadataManager。
     * 此方法确保在关闭Raft节点管理器之前，停止所有调度的任务。
     */
    public void shutdown() {
        // 【新增】关闭协调任务
        if (this.partitionLeaderCoordinator != null) {
            this.partitionLeaderCoordinator.shutdownNow();
        }
        this.raftNodeManager.shutdown();
    }

    /**
     * 在Raft节点成为Leader时调用。
     * 此方法记录当前的Leader任期，并触发一次再均衡检查。
     *
     * @param term 当前的Raft任期。
     */
    @Override
    public void onLeaderStart(long term) {
        this.leaderTerm.set(term);
        log.info("MetadataManager becomes leader at term {}.", term);
        // --- 【优化】: 触发一个事件，通知监听器此节点已成为 Leader ---
        // 这将被 BrokerServer 用于事件驱动的默认凭据管理。
        // 注意：这需要您在 MetadataChangeEvent.EventType 枚举中添加 METADATA_LEADER_ELECTED
        notifyListeners(new MetadataChangeEvent(MetadataChangeEvent.EventType.METADATA_LEADER_ELECTED, Collections.emptyMap()));

        // 成为 Leader 后，立即基于当前的存活节点进行一次再均衡检查。
        // 【修复】: 重用现有的调度器来执行任务，避免创建和泄露一次性线程池。
        // 使用 schedule 给予集群一个短暂的稳定时间。
        // 【最终修复】: 增加空值检查，防止在启动的特定时间点发生 NullPointerException。
        // onLeaderStart 可能会在 partitionLeaderCoordinator 初始化之前被调用。
        if (this.partitionLeaderCoordinator != null) {
            this.partitionLeaderCoordinator.schedule(this::triggerRebalance, 100, TimeUnit.MILLISECONDS);
        } else {
            log.warn("onLeaderStart called, but partitionLeaderCoordinator is not yet initialized. Skipping initial rebalance trigger.");
        }
    }


    /**
     * 当Leader停止时调用，释放当前的Leader任期。
     *
     * @param status 停止的状态信息，表明停止的原因。
     */
    @Override
    public void onLeaderStop(Status status) {
        this.leaderTerm.set(-1);
        log.info("MetadataManager steps down as leader. Status: {}", status);
    }

    /**
     * 检查当前节点是否为Leader。
     *
     * @return 如果当前节点是Leader，则返回true；否则返回false。
     */
    public boolean isLeader() {
        return this.raftNodeManager.isLeader();
    }

    /**
     * 获取当前Leader的地址。
     *
     * @return 当前Leader的地址，如果没有Leader则返回null。
     */
    public String getLeaderAddress() {
        return this.raftNodeManager.getLeaderAddress();
    }

    /**
     * 【新增】一个周期性任务，用于检查并修复分区Leader。
     * 这是保证高可用的核心机制之一。
     */
    private void coordinatePartitionLeaders() {
        // 确保只有Leader才能执行此操作
        if (!isLeader()) {
            return;
        }

        log.debug("开始分区Leader协调检查...");
        List<BrokerId> aliveBrokers = this.raftNodeManager.getAlivePeers().stream()
                .map(peer -> BrokerIdUtils.fromRaftPeerWithBusinessPort(peer, getBusinessPortForHost(peer.getIp(), peer.getPort())))
                .collect(Collectors.toList());

        if (aliveBrokers.isEmpty()) {
            log.warn("协调检查中止：没有发现存活的Broker。");
            return;
        }

        // 遍历所有主题的每一个分区
        this.topics.forEach((topicName, topicMetadata) -> {
            topicMetadata.getPartitionAssignments().forEach((partitionId, assignment) -> {
                BrokerId currentLeader = assignment.getLeader();
                TopicPartition topicPartition = new TopicPartition(topicName, partitionId);

                // 检查：当前Leader是否存活？
                if (currentLeader == null || !aliveBrokers.contains(currentLeader)) {
                    log.warn("分区{}的Leader{}已失效或未分配。正在从存活副本中选举新Leader...",
                            topicPartition, currentLeader);

                    // 从该分区的副本列表中，找到第一个还存活的节点作为新的Leader
                    BrokerId newLeader = assignment.getReplicas().stream()
                            .filter(aliveBrokers::contains)
                            .findFirst()
                            .orElse(null);

                    if (newLeader != null) {
                        log.info("为分区{}选举出新的Leader: {}", topicPartition, newLeader);
                        // 提议一个Raft日志来更新这个分区的Leader
                        this.updatePartitionLeader(topicPartition, newLeader, status -> {
                            if (!status.isOk()) {
                                log.error("为分区{}提议新Leader{}失败: {}", topicPartition, newLeader, status);
                            }
                        });
                    } else {
                        log.error("严重警告：分区{}已无任何存活的副本，无法选举新Leader！副本列表: {}, 存活节点: {}",
                                topicPartition, assignment.getReplicas(), aliveBrokers);
                    }
                }
            });
        });
    }


}